import torch
from torch.utils.data import DataLoader
from torch import nn
from torch.nn import functional as F
from torch import optim
import torchvision
from torchvision import datasets
from torchvision import transforms
from matplotlib import pyplot as plt
from utils.plots_function import *

# Step 1: load dataset
batch_size = 64
transform = transforms.Compose([transforms.ToTensor(),
                                transforms.Normalize((0.1307,), (0.3081,))])

train_dataset = datasets.MNIST(root='../dataset/mnist',
                               train=True,
                               transform=transform,
                               download=True)
test_dataset = datasets.MNIST(root='../dataset/mnist',
                              train=False,
                              transform=transform,
                              download=True)

train_dataloader = DataLoader(dataset=train_dataset,
                              batch_size=batch_size,
                              shuffle=True)
test_dataloader = DataLoader(dataset=test_dataset,
                             batch_size=batch_size,
                             shuffle=False)


# Step 2:create neural network structure
class MnistNet(torch.nn.Module):
    def __init__(self):
        super(MnistNet, self).__init__()
        self.linear_01 = torch.nn.Linear(784, 512)
        self.linear_02 = torch.nn.Linear(512, 256)
        self.linear_03 = torch.nn.Linear(256, 128)
        self.linear_04 = torch.nn.Linear(128, 64)
        self.linear_05 = torch.nn.Linear(64, 10)

    def forward(self, x):
        x = x.view(-1, 784)
        x = F.relu(self.linear_01(x))
        x = F.relu(self.linear_02(x))
        x = F.relu(self.linear_03(x))
        x = F.relu(self.linear_04(x))
        return self.linear_05(x)


# Step 3:define train process

mnistNet = MnistNet()
optimizer = optim.SGD(mnistNet.parameters(), lr=0.001, momentum=0.9)
train_loss = []
for epoch in range(3):
    for batch_id, (x, y) in enumerate(train_dataloader):
        # x: [batch_size, channel, height, width]
        # y: [batch_size]
        x = x.view(x.size(0), 784)
        out = mnistNet(x)
        y_onehot = one_hot(y)
        loss = F.mse_loss(out, y_onehot)
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        train_loss.append(loss.item())

        if batch_id % 10 == 0:
            print(epoch, batch_id, loss.item())
plot_curve(train_loss)
